package com.example.gpu;

import jcuda.*;
import jcuda.driver.*;

import static jcuda.driver.JCudaDriver.*;

/**
 * JCuda 12.6.0 版本的GPU运算示例
 * 这个示例展示了如何使用最新版本的JCuda库在GPU上执行向量加法运算
 */
public class JCuda1260Example {
    
    // CUDA核函数代码 - 执行向量加法
    private static final String CUDA_CODE = 
        "extern \"C\" \n" +
        "__global__ void vectorAdd(const float *a, const float *b, float *c, int n)\n" +
        "{\n" +
        "    int i = blockIdx.x * blockDim.x + threadIdx.x;\n" +
        "    if (i < n) {\n" +
        "        c[i] = a[i] + b[i];\n" +
        "    }\n" +
        "}\n";

    /**
     * 主函数，程序入口点
     */
    public static void main(String[] args) {
        System.out.println("JCuda 12.6.0 GPU运算示例");
        System.out.println("================================");
        
        // 初始化JCuda驱动
        initializeJCuda();
        
        // 向量大小 - 可以根据需要调整
        int vectorSize = 1000000;
        
        // 创建和初始化输入向量
        float[] hostVectorA = new float[vectorSize];
        float[] hostVectorB = new float[vectorSize];
        float[] hostVectorC = new float[vectorSize];
        
        for (int i = 0; i < vectorSize; i++) {
            hostVectorA[i] = (float)i;
            hostVectorB[i] = (float)(vectorSize - i);
        }
        
        try {
            // 演示向量加法运算
            performVectorAddition(hostVectorA, hostVectorB, hostVectorC, vectorSize);
            
            // 验证结果
            verifyResults(hostVectorA, hostVectorB, hostVectorC, vectorSize);
            
        } catch (Exception e) {
            System.err.println("执行GPU计算时出错: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    /**
     * 初始化JCuda驱动和检查GPU设备
     */
    private static void initializeJCuda() {
        try {
            // 初始化CUDA驱动API
            cuInit(0);
            
            // 获取系统中的GPU设备数量
            int[] deviceCount = new int[1];
            cuDeviceGetCount(deviceCount);
            
            if (deviceCount[0] == 0) {
                System.out.println("错误: 未检测到支持CUDA的NVIDIA GPU设备");
                System.out.println("请确保您的系统安装了支持CUDA的NVIDIA显卡和正确的驱动程序");
                System.exit(1);
            }
            
            System.out.println("检测到 " + deviceCount[0] + " 个CUDA设备");
            
            // 打印第一个CUDA设备的信息
            CUdevice device = new CUdevice();
            cuDeviceGet(device, 0);
            
            // 获取设备名称
            byte[] deviceName = new byte[256];
            cuDeviceGetName(deviceName, deviceName.length, device);
            System.out.println("使用GPU设备: " + new String(deviceName).trim());
            
            // 获取设备计算能力
            int[] computeCapabilityMajor = new int[1];
            int[] computeCapabilityMinor = new int[1];
            cuDeviceGetAttribute(computeCapabilityMajor, CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device);
            cuDeviceGetAttribute(computeCapabilityMinor, CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device);
            System.out.println("设备计算能力: " + computeCapabilityMajor[0] + "." + computeCapabilityMinor[0]);
            
        } catch (Exception e) {
            System.err.println("JCuda初始化失败: " + e.getMessage());
            System.err.println("请确保您已安装与JCuda 12.6.0兼容的CUDA Toolkit");
            System.exit(1);
        }
    }
    
    /**
     * 在GPU上执行向量加法运算
     */
    private static void performVectorAddition(float[] a, float[] b, float[] c, int n) {
        CUcontext context = null;
        CUmodule module = null;
        CUfunction function = null;
        CUdeviceptr devicePointerA = null;
        CUdeviceptr devicePointerB = null;
        CUdeviceptr devicePointerC = null;
        
        try {
            // 创建CUDA上下文
            CUdevice device = new CUdevice();
            cuDeviceGet(device, 0);
            context = new CUcontext();
            cuCtxCreate(context, 0, device);
            
            // 创建并加载CUDA模块和函数
            module = new CUmodule();
            function = new CUfunction();
            
            // 加载编译好的PTX文件
            String ptxFileName = "vectorAdd.ptx";
            System.out.println("\n尝试加载PTX文件: " + ptxFileName);
            
            // 检查PTX文件是否存在
            java.io.File ptxFile = new java.io.File(ptxFileName);
            if (!ptxFile.exists()) {
                System.out.println("错误: PTX文件不存在。请先使用以下命令编译CUDA代码:");
                System.out.println("nvcc -ptx vectorAdd.cu -o vectorAdd.ptx");
                System.out.println("执行CPU计算进行演示...");
                
                // 执行CPU计算作为备选
                long startTime = System.currentTimeMillis();
                for (int i = 0; i < n; i++) {
                    c[i] = a[i] + b[i];
                }
                long endTime = System.currentTimeMillis();
                
                System.out.println("CPU计算完成，耗时: " + (endTime - startTime) + " ms");
                return;
            }
            
            // 加载PTX文件
            cuModuleLoad(module, ptxFileName);
            System.out.println("PTX文件加载成功");
            
            // 获取核函数
            cuModuleGetFunction(function, module, "vectorAdd");
            System.out.println("成功获取vectorAdd核函数");
            
            // 1. 分配设备内存
            devicePointerA = new CUdeviceptr();
            devicePointerB = new CUdeviceptr();
            devicePointerC = new CUdeviceptr();
            cuMemAlloc(devicePointerA, n * Sizeof.FLOAT);
            cuMemAlloc(devicePointerB, n * Sizeof.FLOAT);
            cuMemAlloc(devicePointerC, n * Sizeof.FLOAT);
            
            System.out.println("设备内存分配完成");
            
            // 计时开始
            long startTime = System.currentTimeMillis();
            
            // 2. 将数据从主机内存复制到设备内存
            cuMemcpyHtoD(devicePointerA, Pointer.to(a), n * Sizeof.FLOAT);
            cuMemcpyHtoD(devicePointerB, Pointer.to(b), n * Sizeof.FLOAT);
            
            System.out.println("数据从主机复制到设备完成");
            
            // 3. 设置线程块和网格大小
            int blockSize = 256;
            int gridSize = (n + blockSize - 1) / blockSize;
            
            System.out.println("配置线程参数: blockSize=" + blockSize + ", gridSize=" + gridSize);
            
            // 4. 准备核函数参数
            Pointer kernelParameters = Pointer.to(
                Pointer.to(devicePointerA),
                Pointer.to(devicePointerB),
                Pointer.to(devicePointerC),
                Pointer.to(new int[]{n})
            );
            
            // 5. 启动核函数
            cuLaunchKernel(
                function,
                gridSize, 1, 1,    // 网格维度
                blockSize, 1, 1,   // 块维度
                0, null,           // 共享内存大小和流
                kernelParameters, null // 内核参数
            );
            
            System.out.println("GPU核函数启动成功");
            
            // 6. 等待核函数执行完成
            cuCtxSynchronize();
            
            // 7. 将结果从设备内存复制回主机内存
            cuMemcpyDtoH(Pointer.to(c), devicePointerC, n * Sizeof.FLOAT);
            
            // 计时结束
            long endTime = System.currentTimeMillis();
            
            System.out.println("计算完成，耗时: " + (endTime - startTime) + " ms");
            System.out.println("向量大小: " + n);
            
        } catch (Exception e) {
            throw new RuntimeException("向量加法运算失败", e);
        } finally {
            // 释放GPU资源
            if (devicePointerA != null) cuMemFree(devicePointerA);
            if (devicePointerB != null) cuMemFree(devicePointerB);
            if (devicePointerC != null) cuMemFree(devicePointerC);
            if (module != null) cuModuleUnload(module);
            if (context != null) cuCtxDestroy(context);
        }
    }
    
    /**
     * 验证计算结果是否正确
     */
    private static void verifyResults(float[] a, float[] b, float[] c, int n) {
        boolean allCorrect = true;
        int checkCount = Math.min(10, n); // 只检查前10个元素和一些随机元素
        
        // 检查前10个元素
        System.out.println("\n验证计算结果:");
        for (int i = 0; i < checkCount; i++) {
            float expected = a[i] + b[i];
            if (Math.abs(c[i] - expected) > 1e-6) {
                System.out.printf("错误: 索引 %d, 预期值 %.2f, 实际值 %.2f\n", i, expected, c[i]);
                allCorrect = false;
            }
        }
        
        // 检查一些随机位置的元素
        for (int i = 0; i < 5; i++) {
            int randomIndex = (int)(Math.random() * n);
            float expected = a[randomIndex] + b[randomIndex];
            if (Math.abs(c[randomIndex] - expected) > 1e-6) {
                System.out.printf("错误: 随机索引 %d, 预期值 %.2f, 实际值 %.2f\n", randomIndex, expected, c[randomIndex]);
                allCorrect = false;
            }
        }
        
        if (allCorrect) {
            System.out.println("计算结果验证通过！");
        }
    }
}